import { EventSourceParserStream } from 'eventsource-parser/stream';
import type { FetchResponse } from 'ofetch';
import { useState } from 'react';
import type { ChunkChatCompletionResponse } from 'types/dify-response';

export interface FlowNode {
  id: string;
  title: string;
  message: string;
  isStreaming: boolean;
}

export function useDify() {
  const [nodes, setNodes] = useState<FlowNode[]>([]);

  const processStream = async (
    request: () => Promise<
      FetchResponse<ReadableStream<Uint8Array<ArrayBufferLike>>>
    >
  ) => {
    try {
      const res = await request();

      return new Promise((resolve, reject) => {
        if (!res.body) {
          throw new Error('请求体不存在');
        }

        let finalOutput = '';

        res.body
          .pipeThrough(new TextDecoderStream())
          .pipeThrough(new EventSourceParserStream())
          .pipeTo(
            new WritableStream({
              write: (chunk) => {
                const message: ChunkChatCompletionResponse = JSON.parse(
                  chunk.data
                );

                switch (message.event) {
                  case 'message':
                    setNodes((prevNodes) => {
                      return prevNodes.map((node) => {
                        if (message.from_variable_selector?.includes(node.id)) {
                          return {
                            ...node,
                            message: node.message + (message.answer || ''),
                          };
                        }
                        return node;
                      });
                    });
                    break;

                  case 'node_started':
                    if (message.data.node_type === 'llm') {
                      setNodes((prevNodes) => [
                        ...prevNodes,
                        {
                          id: message.data.node_id,
                          title: message.data.title,
                          isStreaming: true,
                          message: '',
                        },
                      ]);
                    }
                    break;

                  case 'node_finished':
                    if (message.data?.outputs?.answer) {
                      finalOutput = message.data.outputs.answer;
                    }

                    setNodes((prevNodes) => {
                      return prevNodes.map((node) => {
                        if (node.id === message.data.node_id) {
                          return {
                            ...node,
                            isStreaming: false,
                          };
                        }

                        return node;
                      });
                    });
                    break;

                  case 'workflow_started':
                    break;

                  case 'workflow_finished':
                    resolve(finalOutput as string);
                    break;
                }
              },
            })
          )
          .catch(reject);
      });
    } catch (error) {
      throw error;
    }
  };

  return {
    nodes,
    processStream,
  };
}
